import torch
from torch import nn, Tensor
import random
import math
from network_files.image_list import ImageList
from torch.jit.annotations import List, Tuple, Dict, Optional
import torchvision


class GeneralizedRCNNTransform(nn.Module):
    """
    Performs input / target transformation before feeding the data to a GeneralizedRCNN
    model.

    The transformations it perform are:
        - input normalization (mean subtraction and std division)
        - input / target resizing to match min_size / max_size

    It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
    """

    def __init__(self, min_size, max_size, image_mean, image_std):
        '''
        __init__方法在类初始化的时候调用的方法
        __init__方法的参数也是在初始化类时需要传入的方法
        通常用在__init__方法中传入一些必要的参数以及执行一些必要的前置代码和初始化代码

        Arguments
            min_size:指定图像的最小边长范围，type -> list or tuple
            max_size:指定图像的最大边长范围，type -> list or tuple
            image_mean:指定图像在标准化处理中的均值，通常是训练数据均值或其他数据集均值   type -> list[3]
            image_std:指定图像在标准化处理中的方差，通常是训练数据方差或其他数据集方差  type -> list[3]
        '''

        ## super()表示是父类的初始化方法，也就是父类的__init__方法
        super(GeneralizedRCNNTransform, self).__init__()
        ## 容错代码。isinstance判断传入参数是否是同一格式，返回True或者False
        if not isinstance(min_size, (list, tuple)):
            min_size = (min_size,)
        self.min_size = min_size      
        self.max_size = max_size      
        self.image_mean = image_mean  
        self.image_std = image_std    

    def normalize(self, image):
        """标准化处理"""
        ## 获取设备信息和数据类型，图像image现在是Tensor格式
        ## 和Tensor进行计算需要是同种类型数据以及在同一设备上，这里的设备可以是cpu也可以是某块GPU
        ## 由于要和图片image进行计算，所以这里获取image的type和device
        dtype, device = image.dtype, image.device
        ## image_mean代表均值，将均值转成Tensor
        mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
        ## image_std表示方差，将方差转成Tensor
        std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
        ## image的shape是[channel,height, width]，但是现在mean和std是[3]，在计算时需要进行下转换
        ## [:, None, None]: shape [3] -> [3, 1, 1]
        return (image - mean[:, None, None]) / std[:, None, None]

    def torch_choice(self, l):
        # type: (List[int])
        """
        Implements `random.choice` via torch ops so it can be compiled with
        TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
        is fixed.
        """
        index = int(torch.empty(1).uniform_(0., float(len(l))).item())
        return l[index]

    def resize(self, image, target):
        # type: (Tensor, Optional[Dict[str, Tensor]])
        """
        将图片缩放到指定的大小范围内，并对应缩放bboxes信息
        Args:
            image: 输入的图片
            target: 输入图片的相关信息（包括bboxes信息）

        Returns:
            image: 缩放后的图片
            target: 缩放bboxes后的图片相关信息
        """
        # 获取输入图像的高和宽
        h, w = image.shape[-2:]
        # 将高和宽值编程Tensor
        im_shape = torch.tensor(image.shape[-2:])
        # 获取高宽中的最小值
        min_size = float(torch.min(im_shape))  
        # 获取高宽中的最大值
        max_size = float(torch.max(im_shape))  
        if self.training:
            # 指定输入图片的最小边长,注意是self.min_size不是min_size
            # torch_choice 实现的功能和random.choice一样，对Tensor进行处理
            size = float(self.torch_choice(self.min_size))  
        else:
            # 指定输入图片的最小边长,注意是self.min_size不是min_size
            size = float(self.min_size[-1])    
        # 根据指定最小边长和图片最小边长计算缩放比例
        scale_factor = size / min_size  

        # 如果使用该缩放比例计算的图片最大边长大于指定的最大边长
        if max_size * scale_factor > self.max_size:
            # 将缩放比例设为指定最大边长和图片最大边长之比
            scale_factor = self.max_size / max_size  

        # interpolate利用插值的方法缩放图片
        # image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]
        # bilinear只支持4D Tensor
        image = torch.nn.functional.interpolate(
            image[None], scale_factor=scale_factor, mode='bilinear', align_corners=False)[0]
        # 如果不是训练模式，target是None。直接返回结果
        if target is None:
            return image, target
        # 训练模式情况下，需要对target中的bbox进行转换，因为bbox的坐标是原图resize之前的坐标
        bbox = target["boxes"]
        # 根据图像的缩放比例来缩放bbox
        bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
        target["boxes"] = bbox

        return image, target

    def max_by_axis(self, the_list):
        # type: (List[List[int]]) -> List[int]
        ## 将第一个赋值给max
        maxes = the_list[0]
        for sublist in the_list[1:]:
            for index, item in enumerate(sublist):
                maxes[index] = max(maxes[index], item)
        return maxes

    def batch_images(self, images, size_divisible=32):
        # type: (List[Tensor], int)
        """
        将一批图像打包成一个batch返回（注意batch中每个tensor的shape是相同的）
        Args:
            images: 输入的一批图片
            size_divisible: 将图像高和宽调整到该数的整数倍

        Returns:
            batched_imgs: 打包成一个batch后的tensor数据
        """

        # 分别计算一个batch中所有图片中的最大channel, height, width
        max_size = self.max_by_axis([list(img.shape) for img in images])
        # 转换成float类型，方便参与计算
        stride = float(size_divisible)
        # max_size = list(max_size)
        # 将height向上调整到stride的整数倍
        max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)
        # 将width向上调整到stride的整数倍
        max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)

        # [batch, channel, height, width]
        batch_shape = [len(images)] + max_size

        # 创建shape为batch_shape且值全部为0的tensor
        batched_imgs = images[0].new_full(batch_shape, 0)
        for img, pad_img in zip(images, batched_imgs):
            # 将输入images中的每张图片复制到新的batched_imgs的每张图片中，对齐左上角，保证bboxes的坐标不变
            # 这样保证输入到网络中一个batch的每张图片的shape相同
            # copy_: Copies the elements from src into self tensor and returns self
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)

        return batched_imgs

    def postprocess(self, result, image_shapes, original_image_sizes):
        # type: (List[Dict[str, Tensor]], List[Tuple[int, int]], List[Tuple[int, int]])
        """
        对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
        Args:
            result: list(dict), 网络的预测结果, len(result) == batch_size
            image_shapes: list(torch.Size), 图像预处理缩放后的尺寸, len(image_shapes) == batch_size
            original_image_sizes: list(torch.Size), 图像的原始尺寸, len(original_image_sizes) == batch_size

        Returns:

        """
        if self.training:
            # 训练过程不需要进行转换，只需要根据标注信息计算loss进行反向传播
            return result
        # 这部分内容没有实质的复杂逻辑。只是遍历图片数据获取原始图片尺寸与缩放后的图片尺寸，调用resize_boxes方法
        for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
            boxes = pred["boxes"]
            boxes = resize_boxes(boxes, im_s, o_im_s)  # 将bboxes缩放回原图像尺度上
            result[i]["boxes"] = boxes
        return result


    def forward(self, images, targets=None):
        # type: (List[Tensor], Optional[List[Dict[str, Tensor]]])
        ## 将images转成List结构
        images = [img for img in images]
        ## 开始遍历数据并进行处理
        for i in range(len(images)):
            image = images[i]
            ## 容错代码，从逻辑上targets可以是None，train过程和inference过程需要对数据做同样的操作
            ## 在train的过程中targets为非None
            ## inference过程并不需要targets
            target_index = targets[i] if targets is not None else None
            ## 容错代码。dim()返回Tensor是几维的，这里验证数据合法性
            if image.dim() != 3:
                raise ValueError("images is expected to be a list of 3d tensors "
                                 "of shape [C, H, W], got {}".format(image.shape))
            ## 对图像进行标准化处理
            image = self.normalize(image)                
            ## 对图像和对应的bboxes缩放到指定范围
            ## 这里并不是把每张图片都强行缩放到一个尺寸，而是将图片按比例缩放到一个范围内，保证图片的原始比例
            image, target_index = self.resize(image, target_index) 
            ## 把经过标准化和resize处理后的数据放回原列表中  
            images[i] = image
            if targets is not None and target_index is not None:
                targets[i] = target_index

        ## 将resize之后的图片宽高保存下来，形成list
        image_sizes = [img.shape[-2:] for img in images]
        ## 将images打包成一个batch
        images = self.batch_images(images)
        ## 定义List[Tuple[int, int]]结构List，初始值为[]
        image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])

        ## 这里将前面收集的图片resize之后的尺寸遍历，放入新创建的image_sizes_list之中
        for image_size in image_sizes:
            ## 容错代码  验证是否有宽和高
            assert len(image_size) == 2
            image_sizes_list.append((image_size[0], image_size[1]))
        
        ## ImageList中携带图像的tensor以及resize之后的图像尺寸
        ## 因为最后输出的是相对于resize之后的图像的框，我们需要转换成原图像的框
        ## 原图尺寸在targets里存储，resize之后的尺寸在ImageList中
        image_list = ImageList(images, image_sizes_list)
        return image_list, targets


def resize_boxes(boxes, original_size, new_size):
    # type: (Tensor, List[int], List[int]) -> Tensor
    """
    将boxes参数根据图像的缩放情况进行相应缩放

    Arguments:
        original_size: 图像缩放前的尺寸
        new_size: 图像缩放后的尺寸
    """
    # 下列列表生成式中，s分别表示缩放后图像高和宽，s_orig分别表示原图像高和宽
    # s与s_orig转换成Tensor之后相除，得到了原图像与缩放后图像高和宽的缩放比例，类型为Tensor
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=boxes.device) /
        torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratios_height, ratios_width = ratios
    # boxes的shape=[mini_batch, 4]
    # unbind在dim=1的维度进行拆分，xmin,ymin,xmax,ymax分别是[mini_batch]
    xmin, ymin, xmax, ymax = boxes.unbind(1)
    # x y坐标分别乘以对应的缩放比例得到缩放后的x y坐标
    xmin = xmin * ratios_width
    xmax = xmax * ratios_width
    ymin = ymin * ratios_height
    ymax = ymax * ratios_height
    return torch.stack((xmin, ymin, xmax, ymax), dim=1)








